import pandas as pd
from datetime import datetime
from Util.Graph import Graph
from Util.Edge import Edge
from tqdm import tqdm


class DataLoader:

    def __init__(self, end, name='HepTh'):
        start_time = 0

        if name == 'HepTh':
            path = "../Data/ca-cit-HepTh/out.ca-cit-HepTh"
            start_time = datetime.fromtimestamp(0)
            end_time = datetime.fromtimestamp(end)
        self.end_time = end_time
        self.start_time = start_time
        self.df = pd.read_table(path, sep=' ', parse_dates=[3], header=None, skiprows=2,
                                date_parser=lambda dates: datetime.fromtimestamp(int(dates)), error_bad_lines=False)
        self.graph = Graph()

    def read_data(self):
        min_year = 1e30
        lines = self.df.values.tolist()
        for line in tqdm(lines):
            start = line[0]
            end = line[1]
            time = line[3]
            weight = float(line[2])
            if (time.year > 1970) & (time < self.end_time):
                if time.year < min_year:
                    min_year = time.year
                edge = Edge(start, end, time, weight)
                self.graph.add_edge(edge)
        self.graph.start_time = min_year
